import os
import numpy as np
import matplotlib.pyplot as plt
import argparse

from vehicle_model import DoubleIntegratorVehicle
from lqr_controller import LQRController
from trajectory_generator import TrajectoryGenerator
from simulation import Simulation
from visualization import Visualization
from utils import calculate_errors, set_plot_style

def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='双积分小车LQR轨迹跟踪控制')
    parser.add_argument('--dt', type=float, default=0.1, help='仿真步长 (默认: 0.1)')
    parser.add_argument('--sim-time', type=float, default=30.0, help='仿真时间 (默认: 30.0)')
    parser.add_argument('--q-pos', type=float, default=2000.0, help='位置状态权重 (默认: 2000.0)')
    parser.add_argument('--q-vel', type=float, default=200.0, help='速度状态权重 (默认: 200.0)')
    parser.add_argument('--r-acc', type=float, default=0.01, help='加速度控制权重 (默认: 0.01)')
    parser.add_argument('--trajectory', type=str, default='circle', 
                        choices=['circle', 'figure_eight', 'square'], 
                        help='轨迹类型 (默认: circle)')
    parser.add_argument('--output-dir', type=str, default='results', 
                        help='输出目录 (默认: results)')
    parser.add_argument('--no-animation', action='store_true', help='不生成动画')
    parser.add_argument('--debug', action='store_true', help='调试模式')
    return parser.parse_args()

def create_output_dir(output_dir):
    """创建输出目录"""
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        print(f"创建输出目录: {output_dir}")

def main():
    """主函数"""
    # 解析命令行参数
    args = parse_arguments()
    
    # 创建输出目录
    create_output_dir(args.output_dir)
    
    # 设置绘图样式
    set_plot_style()
    
    # 创建系统组件
    vehicle = DoubleIntegratorVehicle(dt=args.dt)
    traj_gen = TrajectoryGenerator(dt=args.dt)
    
    # 获取系统矩阵
    A, B = vehicle.get_state_space_matrices()
    
    # 创建LQR控制器
    Q = np.diag([args.q_pos, args.q_pos, args.q_vel, args.q_vel])
    R = np.diag([args.r_acc, args.r_acc])
    controller = LQRController(A, B, Q, R)
    
    # 创建仿真器
    simulator = Simulation(vehicle, controller, traj_gen)
    
    # 创建可视化工具
    visualizer = Visualization()
    
    # 准备仿真
    time_points = np.arange(0, args.sim_time, args.dt)
    
    # 生成参考轨迹
    reference_trajectory = traj_gen.generate_trajectory(args.trajectory, time_points)
    
    # 初始状态设置为轨迹起点
    initial_state = reference_trajectory[0].copy()
    
    # 打印仿真配置
    print("=" * 50)
    print("双积分小车LQR轨迹跟踪控制仿真")
    print("=" * 50)
    print(f"仿真步长: {args.dt}")
    print(f"仿真时间: {args.sim_time}")
    print(f"轨迹类型: {args.trajectory}")
    print(f"参考轨迹点数: {len(reference_trajectory)}")
    print(f"初始状态: {initial_state}")
    print("LQR参数:")
    print(f"  Q (位置权重): {args.q_pos}")
    print(f"  Q (速度权重): {args.q_vel}")
    print(f"  R (加速度权重): {args.r_acc}")
    print(f"输出目录: {args.output_dir}")
    print("=" * 50)
    
    # 首先单独绘制参考轨迹
    plt.figure(figsize=(10, 8))
    plt.plot(reference_trajectory[:, 0], reference_trajectory[:, 1], 'b-', linewidth=3, label='参考轨迹')
    plt.scatter(reference_trajectory[0, 0], reference_trajectory[0, 1], c='g', s=100, marker='o', label='起点')
    plt.scatter(reference_trajectory[-1, 0], reference_trajectory[-1, 1], c='r', s=100, marker='x', label='终点')
    
    # 绘制方向箭头
    step = len(reference_trajectory) // 20  # 约20个箭头
    for i in range(0, len(reference_trajectory), step):
        if i + 5 < len(reference_trajectory):
            plt.arrow(reference_trajectory[i, 0], reference_trajectory[i, 1],
                     (reference_trajectory[i+5, 0] - reference_trajectory[i, 0])/2,
                     (reference_trajectory[i+5, 1] - reference_trajectory[i, 1])/2,
                     head_width=0.4, head_length=0.6, fc='blue', ec='blue', alpha=0.6)
    
    plt.title(f'{args.trajectory}参考轨迹', fontsize=16)
    plt.xlabel('X位置', fontsize=14)
    plt.ylabel('Y位置', fontsize=14)
    plt.grid(True)
    plt.axis('equal')
    plt.legend(fontsize=14)
    plt.savefig(f'{args.output_dir}/reference_trajectory.png', dpi=300, bbox_inches='tight')
    plt.show()
    
    # 运行仿真
    print("\n开始仿真...")
    results = simulator.run_simulation(reference_trajectory, initial_state, time_points, debug=args.debug)
    print("仿真完成!")
    
    # 计算误差
    errors = calculate_errors(reference_trajectory[:len(results['states_history'])], results['states_history'])
    print(f"RMSE: {errors['rmse']:.4f}")
    print(f"最大误差: {errors['max_error']:.4f}")
    
    # 可视化结果
    output_path = f"{args.output_dir}/{args.trajectory}_tracking.png"
    visualizer.plot_trajectory(results, errors, output_path)
    
    # 生成动画
    if not args.no_animation:
        animation_path = f"{args.output_dir}/{args.trajectory}_animation.gif"
        visualizer.create_animation(results, animation_path)
    
    print("\n输出文件:")
    print(f"- 轨迹图: {output_path}")
    if not args.no_animation:
        print(f"- 动画: {animation_path}")
    
    # 返回结果和误差以便可能的后续分析
    return results, errors

if __name__ == "__main__":
    main() 